package zy.learn.demo.structuredstreaming.watermark

import java.sql.Timestamp

import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.streaming.OutputMode

/**
 * 滚动窗口，使用水印，Append模式，获取window.start
 */
object WatermarkAppend {
  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().set("spark.sql.shuffle.partitions", "3")
    val spark = SparkSession.builder()
      .master("local[2]")
      .config(sparkConf)
      .appName("Watermark Tumbling Window Append")
      .getOrCreate()

    import spark.implicits._

    val lines = spark.readStream
      .format("socket") // 设置数据源
      .option("host", "co7-203")
      .option("port", 9999)
      .load
    /* 输入的数据
    * 2020-10-14 10:55:00,dog     # wartermark: 10:55:00 - 2 min = 10:53:00
    * 2020-10-14 11:00:00,dog     # wartermark: 11:00:00 - 2 min = 10:58:00
    * 2020-10-14 10:55:00,dog     # wartermark: 10:55:00 - 2 min = 10:53:00 < 10:58:00 => 10:58:00
    * 2020-10-14 11:05:00,dog     # wartermark: 11:05:00 - 2 min = 11:03:00
    * 2020-10-14 10:55:00,dog     # wartermark: 11:03:00
    * 2020-10-14 11:10:00,dog     # wartermark: 11:10:00 - 2 min = 11:08:00
    * 2020-10-14 11:20:00,dog     # wartermark: 11:20:00 - 2 min = 11:18:00
    * 2020-10-14 11:30:00,dog     # wartermark: 11:30:00 - 2 min = 11:28:00
    * 输出的结果
+-------------------+-------------------+----+-----+
|start              |end                |word|count|
+-------------------+-------------------+----+-----+
                                                      Batch: 0
+-------------------+-------------------+----+-----+
                                                      Batch: 1
+-------------------+-------------------+----+-----+
                                                      Batch: 2
+-------------------+-------------------+----+-----+
                                                      Batch: 3
+-------------------+-------------------+----+-----+
|2020-10-14 10:50:00|2020-10-14 11:00:00|dog |2    |  Batch: 4
+-------------------+-------------------+----+-----+
                                                      Batch: 5
+-------------------+-------------------+----+-----+
                                                      Batch: 6
+-------------------+-------------------+----+-----+
|2020-10-14 11:00:00|2020-10-14 11:10:00|dog |2    |  Batch: 7
+-------------------+-------------------+----+-----+
+-----+---+----+-----+    Batch: 4
+-----+---+----+-----+
    * */
    val wordsDF = lines.as[String].map(line => {
      val split = line.split(",")
      (Timestamp.valueOf(split(0)), split(1))
    }).toDF("ts", "word")

    import org.apache.spark.sql.functions._

    val wordCounts = wordsDF
      .withWatermark("ts", "2 minutes")
      .groupBy(
        // 调用 window 函数, 返回的是一个 Column 参数 1: df 中表示时间戳的列 参数 2: 窗口长度 参数 3: 滑动步长
        window($"ts", "10 minutes"),
        $"word"
      ).count().select("window.start", "window.end", "word", "count") // 计数

    val query = wordCounts.writeStream
      .format("console")
      .outputMode(OutputMode.Append())
      .option("truncate", "false")
      .start()

    query.awaitTermination()
  }
}
